#!/user/bin/env python3
# -*- coding: utf-8 -*-
"""
--------------------------------------
    Author:     JiChao_Song
    Date  :     2022/2/15 17:05
    Desc  :
--------------------------------------
"""
import json
from argparse import Namespace
import inspect
import logging
import os
import importlib
import signal
import socket
import sys
import time
import atexit
from html import escape
from itertools import chain

import gevent
from locust import stats, log
from locust.argument_parser import get_parser
from locust.runners import MasterRunner, Runner
from locust.util.rounding import proper_round
import locust

from locust.env import Environment
from locust.log import setup_logging, greenlet_exception_logger
from locust.stats import print_error_report, print_percentile_stats, print_stats, stats_printer, stats_history, \
    CONSOLE_STATS_INTERVAL_SEC, sort_stats, console_logger
from locust.stats import StatsCSV, StatsCSVFileWriter
from locust.user import User
from locust.user.inspectuser import print_task_ratio, print_task_ratio_json
from locust.util.timespan import parse_timespan
from locust.exception import AuthCredentialsError
from locust.shape import LoadTestShape
from locust.input_events import input_listener
from locust.html import get_html_report
from json import dumps

from extends.flask_locust.user import PerformanceUser

version = locust.__version__


def is_user_class(item):
    """
    Check if a variable is a runnable (non-abstract) User class
    """
    return bool(inspect.isclass(item) and issubclass(item, User) and item.abstract is False)


def is_shape_class(item):
    """
    Check if a class is a LoadTestShape
    """
    return bool(
        inspect.isclass(item) and issubclass(item, LoadTestShape) and item.__dict__["__module__"] != "locust.shape"
    )


def load_locustfile(path):
    """
    Import given locustfile path and return (docstring, callables).

    Specifically, the locustfile's ``__doc__`` attribute (a string) and a
    dictionary of ``{'name': callable}`` containing all callables which pass
    the "is a Locust" test.
    """

    # Start with making sure the current working dir is in the sys.path
    sys.path.insert(0, os.getcwd())
    # Get directory and locustfile name
    directory, locustfile = os.path.split(path)
    # If the directory isn't in the PYTHONPATH, add it so our import will work
    added_to_path = False
    index = None
    if directory not in sys.path:
        sys.path.insert(0, directory)
        added_to_path = True
    # If the directory IS in the PYTHONPATH, move it to the front temporarily,
    # otherwise other locustfiles -- like Locusts's own -- may scoop the intended
    # one.
    else:
        i = sys.path.index(directory)
        if i != 0:
            # Store index for later restoration
            index = i
            # Add to front, then remove from original position
            sys.path.insert(0, directory)
            del sys.path[i + 1]
    # Perform the import
    source = importlib.machinery.SourceFileLoader(os.path.splitext(locustfile)[0], path)
    imported = source.load_module()
    # Remove directory from path if we added it ourselves (just to be neat)
    if added_to_path:
        del sys.path[0]
    # Put back in original index if we moved it
    if index is not None:
        sys.path.insert(index + 1, directory)
        del sys.path[0]
    # Return our two-tuple
    user_classes = {name: value for name, value in vars(imported).items() if is_user_class(value)}

    # Find shape class, if any, return it
    shape_classes = [value for name, value in vars(imported).items() if is_shape_class(value)]
    if shape_classes:
        shape_class = shape_classes[0]()
    else:
        shape_class = None

    return imported.__doc__, user_classes, shape_class


def create_environment(user_classes, options, events=None, shape_class=None, locustfile=None):
    """
    Create an Environment instance from options
    """
    return Environment(
        locustfile = locustfile,
        user_classes = user_classes,
        shape_class = shape_class,
        events = events,
        host = options.host,
        reset_stats = options.reset_stats,
        stop_timeout = options.stop_timeout,
        parsed_options = options,
    )


def stats_requests(runner: Runner):
    stats = []

    for s in chain(sort_stats(runner.stats.entries), [runner.stats.total]):
        stats.append(
            {
                "method": s.method,
                "name": s.name,
                "safe_name": escape(s.name, quote = False),
                "num_requests": s.num_requests,
                "num_failures": s.num_failures,
                "avg_response_time": s.avg_response_time,
                "min_response_time": 0 if s.min_response_time is None else proper_round(s.min_response_time),
                "max_response_time": proper_round(s.max_response_time),
                "current_rps": s.current_rps,
                "current_fail_per_sec": s.current_fail_per_sec,
                "median_response_time": s.median_response_time,
                "ninetieth_response_time": s.get_response_time_percentile(0.9),
                "ninety_ninth_response_time": s.get_response_time_percentile(0.99),
                "avg_content_length": s.avg_content_length,
            }
        )

    errors = []
    for e in runner.errors.values():
        err_dict = e.to_dict()
        err_dict["name"] = escape(err_dict["name"])
        err_dict["error"] = escape(err_dict["error"])
        errors.append(err_dict)

    # Truncate the total number of stats and errors displayed since a large number of rows will cause the app
    # to render extremely slowly. Aggregate stats should be preserved.
    report = {"stats": stats[:500], "errors": errors[:500]}
    if len(stats) > 500:
        report["stats"] += [stats[-1]]

    if stats:
        report["total_rps"] = stats[len(stats) - 1]["current_rps"]
        report["fail_ratio"] = runner.stats.total.fail_ratio
        report[
            "current_response_time_percentile_95"
        ] = runner.stats.total.get_current_response_time_percentile(0.95)
        report[
            "current_response_time_percentile_50"
        ] = runner.stats.total.get_current_response_time_percentile(0.5)

    is_distributed = isinstance(runner, MasterRunner)
    if is_distributed:
        workers = []
        for worker in runner.clients.values():
            workers.append(
                {
                    "id": worker.id,
                    "state": worker.state,
                    "user_count": worker.user_count,
                    "cpu_usage": worker.cpu_usage,
                    "memory_usage": worker.memory_usage,
                }
            )

        report["workers"] = workers

    report["state"] = runner.state
    report["user_count"] = runner.user_count

    # return jsonify(report)
    # 保存数据

    print(json.dumps(report, ensure_ascii = False))


def save_stats_entries(runner: Runner):
    def stats_printer_func():
        while True:
            # stats = []
            # for s in chain(sort_stats(stats_dict.entries), [stats_dict.total]):
            #     stats.append(
            #         {
            #             "method": s.method,
            #             "name": s.name,
            #             "safe_name": escape(s.name, quote=False),
            #             "num_requests": s.num_requests,
            #             "num_failures": s.num_failures,
            #             "avg_response_time": s.avg_response_time,
            #             "min_response_time": 0 if s.min_response_time is None else proper_round(s.min_response_time),
            #             "max_response_time": proper_round(s.max_response_time),
            #             "current_rps": s.current_rps,
            #             "current_fail_per_sec": s.current_fail_per_sec,
            #             "median_response_time": s.median_response_time,
            #             "ninetieth_response_time": s.get_response_time_percentile(0.9),
            #             "ninety_ninth_response_time": s.get_response_time_percentile(0.99),
            #             "avg_content_length": s.avg_content_length,
            #         }
            #     )
            # print(stats)
            stats_requests(runner)

            gevent.sleep(CONSOLE_STATS_INTERVAL_SEC)

    return stats_printer_func


def main():
    # find specified locustfile and make sure it exists, using a very simplified
    # command line parser that is only used to parse the -f option

    # import the locustfile
    docstring, user_classes, shape_class = None, {'name': PerformanceUser}, None

    # parse all command line options
    options = get_parser().parse_args()
    options.headless = True
    options.spawn_rate = 1
    options.num_users = 10

    if options.headful:
        options.headless = False

    if options.slave or options.expect_slaves:
        sys.stderr.write("The --slave/--expect-slaves parameters have been renamed --worker/--expect-workers\n")
        sys.exit(1)

    if options.autoquit != -1 and not options.autostart:
        sys.stderr.write("--autoquit is only meaningful in combination with --autostart\n")
        sys.exit(1)

    if options.step_time or options.step_load or options.step_users or options.step_clients:
        sys.stderr.write(
            "The step load feature was removed in Locust 1.3. You can achieve similar results using a LoadTestShape class. See https://docs.locust.io/en/stable/custom-load-shape.html\n"
        )
        sys.exit(1)

    if options.hatch_rate:
        sys.stderr.write("[DEPRECATED] The --hatch-rate parameter has been renamed --spawn-rate\n")
        options.spawn_rate = options.hatch_rate

    # setup logging
    if not options.skip_log_setup:
        if options.loglevel.upper() in ["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"]:
            setup_logging(options.loglevel, options.logfile)
        else:
            sys.stderr.write("Invalid --loglevel. Valid values are: DEBUG/INFO/WARNING/ERROR/CRITICAL\n")
            sys.exit(1)

    logger = logging.getLogger(__name__)
    greenlet_exception_handler = greenlet_exception_logger(logger)

    if options.list_commands:
        print("Available Users:")
        for name in user_classes:
            print("    " + name)
        sys.exit(0)

    if not user_classes:
        logger.error("No User class found!")
        sys.exit(1)

    # make sure specified User exists
    if options.user_classes:
        missing = set(options.user_classes) - set(user_classes.keys())
        if missing:
            logger.error("Unknown User(s): %s\n" % (", ".join(missing)))
            sys.exit(1)
        else:
            names = set(options.user_classes) & set(user_classes.keys())
            user_classes = [user_classes[n] for n in names]
    else:
        # list() call is needed to consume the dict_view object in Python 3
        user_classes = list(user_classes.values())

    if os.name != "nt" and not options.master:

        try:
            import resource

            minimum_open_file_limit = 10000
            current_open_file_limit = resource.getrlimit(resource.RLIMIT_NOFILE)[0]

            if current_open_file_limit < minimum_open_file_limit:
                # Increasing the limit to 10000 within a running process should work on at least MacOS.
                # It does not work on all OS:es, but we should be no worse off for trying.
                resource.setrlimit(resource.RLIMIT_NOFILE, [minimum_open_file_limit, resource.RLIM_INFINITY])
        except BaseException:
            logger.warning(
                (
                    f"System open file limit '{current_open_file_limit}' is below minimum setting '{minimum_open_file_limit}'. "
                    "It's not high enough for load testing, and the OS didn't allow locust to increase it by itself. "
                    "See https://github.com/locustio/locust/wiki/Installation#increasing-maximum-number-of-open-files-limit for more info."
                )
            )

    # create locust Environment
    environment = create_environment(
        user_classes, options, events = locust.events, shape_class = shape_class,
        locustfile = os.path.basename('locustfile')
    )

    if shape_class and (options.num_users or options.spawn_rate):
        logger.warning(
            "The specified locustfile contains a shape class but a conflicting argument was specified: users or spawn-rate. Ignoring arguments"
        )

    if options.show_task_ratio:
        print("\n Task ratio per User class")
        print("-" * 80)
        print_task_ratio(user_classes, options.num_users, False)
        print("\n Total task ratio")
        print("-" * 80)
        print_task_ratio(user_classes, options.num_users, True)
        sys.exit(0)
    if options.show_task_ratio_json:
        print_task_ratio_json(user_classes, options.num_users)
        sys.exit(0)

    if options.master:
        if options.worker:
            logger.error("The --master argument cannot be combined with --worker")
            sys.exit(-1)
        runner = environment.create_master_runner(
            master_bind_host = options.master_bind_host,
            master_bind_port = options.master_bind_port,
        )
    elif options.worker:
        try:
            runner = environment.create_worker_runner(options.master_host, options.master_port)
            logger.debug("Connected to locust master: %s:%s", options.master_host, options.master_port)
        except socket.error as e:
            logger.error("Failed to connect to the Locust master: %s", e)
            sys.exit(-1)
    else:
        runner = environment.create_local_runner()

    # main_greenlet is pointing to runners.greenlet by default, it will point the web greenlet later if in web mode
    main_greenlet = runner.greenlet

    if options.run_time:
        if options.worker:
            logger.error("--run-time should be specified on the master node, and not on worker nodes")
            sys.exit(1)
        try:
            options.run_time = parse_timespan(options.run_time)
        except ValueError:
            logger.error("Valid --run-time formats are: 20, 20s, 3m, 2h, 1h20m, 3h30m10s, etc.")
            sys.exit(1)

    if options.csv_prefix:
        stats_csv_writer = StatsCSVFileWriter(
            environment, stats.PERCENTILES_TO_REPORT, options.csv_prefix, options.stats_history_enabled
        )
    else:
        stats_csv_writer = StatsCSV(environment, stats.PERCENTILES_TO_REPORT)

    # start Web UI
    if not options.headless and not options.worker:
        # spawn web greenlet
        protocol = "https" if options.tls_cert and options.tls_key else "http"
        try:
            if options.web_host == "*":
                # special check for "*" so that we're consistent with --master-bind-host
                web_host = ""
            else:
                web_host = options.web_host
            if web_host:
                logger.info("Starting web interface at %s://%s:%s" % (protocol, web_host, options.web_port))
            else:
                logger.info(
                    "Starting web interface at %s://0.0.0.0:%s (accepting connections from all network interfaces)"
                    % (protocol, options.web_port)
                )
            web_ui = environment.create_web_ui(
                host = web_host,
                port = options.web_port,
                auth_credentials = options.web_auth,
                tls_cert = options.tls_cert,
                tls_key = options.tls_key,
                stats_csv_writer = stats_csv_writer,
                delayed_start = True,
            )
        except AuthCredentialsError:
            logger.error("Credentials supplied with --web-auth should have the format: username:password")
            sys.exit(1)
    else:
        if options.autostart:
            logger.warning("Option --autostart is ignored for headless mode and worker process.")
        web_ui = None

    def assign_equal_weights(environment, **kwargs):
        environment.assign_equal_weights()

    if options.equal_weights:
        environment.events.init.add_listener(assign_equal_weights)

    # Fire locust init event which can be used by end-users' code to run setup code that
    # need access to the Environment, Runner or WebUI.
    environment.events.init.fire(environment = environment, runner = runner, web_ui = web_ui)

    if web_ui:
        web_ui.start()
        main_greenlet = web_ui.greenlet

    def stop_and_optionally_quit():
        if options.autostart:
            logger.info("--run-time limit reached, stopping test")
            runner.stop()
            if options.autoquit != -1:
                logger.debug("Autoquit time limit set to %s seconds" % options.autoquit)
                time.sleep(options.autoquit)
                logger.info("--autoquit time reached, shutting down")
                runner.quit()
                if web_ui:
                    web_ui.stop()
            else:
                logger.info("--autoquit not specified, leaving web ui running indefinitely")
        else:  # --headless run
            logger.info("--run-time limit reached. Stopping Locust")
            runner.quit()

    def spawn_run_time_quit_greenlet():
        gevent.spawn_later(options.run_time, stop_and_optionally_quit).link_exception(greenlet_exception_handler)

    headless_master_greenlet = None
    stats_printer_greenlet = None
    if not options.only_summary and (options.print_stats or (options.headless and not options.worker)):
        # spawn stats printing greenlet
        stats_printer_greenlet = gevent.spawn(save_stats_entries(runner))
        stats_printer_greenlet.link_exception(greenlet_exception_handler)

        # stats_printer_greenlet_1 = gevent.spawn(stats_printer(runner.stats))
        # stats_printer_greenlet_1.link_exception(greenlet_exception_handler)

    gevent.spawn(stats_history, runner)

    def start_automatic_run():
        if options.master:
            # wait for worker nodes to connect
            start_time = time.monotonic()
            while len(runner.clients.ready) < options.expect_workers:
                if options.expect_workers_max_wait and options.expect_workers_max_wait < time.monotonic() - start_time:
                    logger.error("Gave up waiting for workers to connect.")
                    runner.quit()
                    sys.exit(1)
                logging.info(
                    "Waiting for workers to be ready, %s of %s connected",
                    len(runner.clients.ready),
                    options.expect_workers,
                )

                # TODO: Handle KeyboardInterrupt and send quit signal to workers that are started.
                #       Right now, if the user sends a ctrl+c, the master will not gracefully
                #       shutdown resulting in all the already started workers to stay active.
                time.sleep(1)
        if not options.worker:
            # apply headless mode defaults
            if options.num_users is None:
                options.num_users = 1
            if options.spawn_rate is None:
                options.spawn_rate = 1

            # start the test
            if environment.shape_class:
                if options.run_time:
                    sys.stderr.write("It makes no sense to combine --run-time and LoadShapes. Bailing out.\n")
                    sys.exit(1)
                environment.runner.start_shape()
                environment.runner.shape_greenlet.join()
                stop_and_optionally_quit()
            else:
                headless_master_greenlet = gevent.spawn(runner.start, options.num_users, options.spawn_rate)
                headless_master_greenlet.link_exception(greenlet_exception_handler)

        if options.run_time:
            logger.info("Run time limit set to %s seconds" % options.run_time)
            spawn_run_time_quit_greenlet()
        elif not options.worker and not environment.shape_class:
            logger.info("No run time limit set, use CTRL+C to interrupt")

    if options.headless:
        start_automatic_run()

    input_listener_greenlet = None
    if not options.worker:
        # spawn input listener greenlet
        input_listener_greenlet = gevent.spawn(
            input_listener(
                {
                    "w": lambda: runner.start(runner.user_count + 1, 100)
                    if runner.state != "spawning"
                    else logging.warning("Already spawning users, can't spawn more right now"),
                    "W": lambda: runner.start(runner.user_count + 10, 100)
                    if runner.state != "spawning"
                    else logging.warning("Already spawning users, can't spawn more right now"),
                    "s": lambda: runner.start(max(0, runner.user_count - 1), 100)
                    if runner.state != "spawning"
                    else logging.warning("Spawning users, can't stop right now"),
                    "S": lambda: runner.start(max(0, runner.user_count - 10), 100)
                    if runner.state != "spawning"
                    else logging.warning("Spawning users, can't stop right now"),
                },
            )
        )
        input_listener_greenlet.link_exception(greenlet_exception_handler)
        # ensure terminal is reset, even if there is an unhandled exception in locust or someone
        # does something wild, like calling sys.exit() in the locustfile
        atexit.register(input_listener_greenlet.kill, block = True)

    if options.csv_prefix:
        gevent.spawn(stats_csv_writer.stats_writer).link_exception(greenlet_exception_handler)

    def shutdown():
        """
        Shut down locust by firing quitting event, printing/writing stats and exiting
        """
        logger.debug("Running teardowns...")

        if input_listener_greenlet is not None:
            input_listener_greenlet.kill(block = False)

        environment.events.quitting.fire(environment = environment, reverse = True)

        # determine the process exit code
        if log.unhandled_greenlet_exception:
            code = 2
        elif environment.process_exit_code is not None:
            code = environment.process_exit_code
        elif len(runner.errors) or len(runner.exceptions):
            code = options.exit_code_on_error
        else:
            code = 0

        logger.info("Shutting down (exit code %s)" % code)
        if stats_printer_greenlet is not None:
            stats_printer_greenlet.kill(block = False)
        if headless_master_greenlet is not None:
            headless_master_greenlet.kill(block = False)
        logger.debug("Cleaning up runner...")
        if runner is not None:
            runner.quit()

        if not isinstance(runner, locust.runners.WorkerRunner):
            print_stats(runner.stats, current = False)
            print_percentile_stats(runner.stats)
            print_error_report(runner.stats)

        sys.exit(code)

    # install SIGTERM handler
    def sig_term_handler():
        logger.info("Got SIGTERM signal")
        shutdown()

    def save_html_report():
        html_report = get_html_report(environment, show_download_link = False)
        logger.info("writing html report to file: %s", options.html_file)
        with open(options.html_file, "w", encoding = "utf-8") as file:
            file.write(html_report)

    gevent.signal_handler(signal.SIGTERM, sig_term_handler)

    try:
        logger.info("Starting Locust %s" % version)
        if options.autostart:
            start_automatic_run()

        main_greenlet.join()
        if options.html_file:
            save_html_report()
    except KeyboardInterrupt:
        if options.html_file:
            save_html_report()
    except Exception:
        raise
    shutdown()
